import zmq, time, cv2
import threading
import numpy as np
from paddle.inference import Config, create_predictor
import json


class lane:
    def __init__(self):
        self.stop = False
        self.server = self.get_server(5001)
        self.flag_infer_initok = False
        model_path = 'paddle_jetson/lane_model/cnn_lane.pdmodel'
        params_path = 'paddle_jetson/lane_model/cnn_lane.pdiparams'
        config = Config(model_path, params_path)
        self.predictor = create_predictor(config)
        thread_tmp = threading.Thread(target=self.process_demo)
        thread_tmp.daemon = True
        thread_tmp.start()

    def init(self):
        img = np.zeros((128, 128, 3), np.uint8)
        infer_tmp = self.crusie
        infer_tmp(img)
        self.flag_infer_initok = True
        while True:
            try:
                if self.stop:
                    break
                time.sleep(1)
            except Exception as e:
                print(e)
                break

    def crusie(self, image):
        img = image.astype(np.float32) / 127.5
        mean = np.array([1.0, 1.0, 1.0])
        if mean is not None:
            img -= mean
        img = img[:, :, ::-1].astype('float32')  # bgr -> rgb
        img = img.transpose((2, 0, 1))  # hwc -> chw
        img = img[np.newaxis, :]

        input_names = self.predictor.get_input_names()
        input_tensor = self.predictor.get_input_handle(input_names[0])
        input_tensor.copy_from_cpu(img)
        self.predictor.run()
        output_names = self.predictor.get_output_names()
        output_tensor = self.predictor.get_output_handle(output_names[0])
        output_data = output_tensor.copy_to_cpu()[0].tolist()
        return output_data

    def get_server(self, port):
        context = zmq.Context()
        socket = context.socket(zmq.REP)
        socket.bind(f"tcp://127.0.0.1:{port}")
        return socket

    def process_demo(self):
        print(time.strftime("%Y-%m-%d %H:%M:%S"), "{} process start".format(self.server))
        while True:
            response = self.server.recv()
            head = response[:5]
            res = []
            if head == b"Astop":
                self.stop = True
                break
            if head == b"ATATA":
                if self.flag_infer_initok:
                    res = True
                else:
                    res = False
            elif head == b"image":
                # 把bytes转为jpg格式
                img = cv2.imdecode(np.frombuffer(response[5:], dtype=np.uint8), 1)
                if self.flag_infer_initok:
                    res = self.crusie(img)
            json_data = json.dumps(res)
            json_data = bytes(json_data, encoding='utf-8')
            self.server.send(json_data)


main = lane()
main.init()
